import argparse

def args_parser():
    parser = argparse.ArgumentParser()

    # federated arguments
    parser.add_argument('--epochs',type=int,default=1000,help="rounds of training")
    parser.add_argument('--num_users',type=int,default=100,help="number of users: K")
    parser.add_argument('--frac', type=float, default=0.10, help="the fraction of clients: C")
    parser.add_argument('--local_ep', type=int, default=5, help="the number of local epochs: E")
    parser.add_argument('--local_bs', type=int, default=50, help="local batch size: B")
    parser.add_argument('--bs', type=int, default=128, help="test batch size")
    parser.add_argument('--algorithm', type=str, default='FedDery', help="FedTest")
    parser.add_argument('--model', type=str, default='cnn_mobilenetv2_resnet18', help="model name")
    parser.add_argument('--device', type=int, default=0, help="GPU ID, -1 for CPU")
    #optimizer
    parser.add_argument('--optimizer', type=str, default='sgd', help='the optimizer')
    parser.add_argument('--lr', type=float, default=0.01, help="learning rate")
    parser.add_argument('--lr_decay', type=float, default=0.998, help="learning rate decay")
    parser.add_argument('--momentum', type=float, default=0.5, help="SGD momentum (default: 0.5)")
    parser.add_argument('--weight_decay', type=float, default=1e-3, help="weight_decay (default: 1e-4)")

    # dataset
    parser.add_argument('--dataset',type=str,default='cifar10',help="name of dataset")
    parser.add_argument('--generate_data', type=int, default=1, help="whether generate new dataset")
    parser.add_argument('--iid', type=int, default=1, help='whether i.i.d or not')
    parser.add_argument('--noniid_case', type=int, default=5, help="non i.i.d case (1, 2, 3, 4)")
    parser.add_argument('--data_beta', type=float, default=0.5,
                        help='The parameter for the dirichlet distribution for data partitioning')
    parser.add_argument('--num_classes', type=int, default=10, help="number of classes")
    parser.add_argument('--num_channels', type=int, default=3, help="number of channels")
    parser.add_argument('--glo_dataset', type=int, default=1, help="whether use the global dataset")
    parser.add_argument('--glo_dataset_frac', type=float, default=0.00, help="the fraction of global dataset")
    parser.add_argument('--seed', type=int, default=1, help="number of seed")
    # FedTest
    # parser.add_argument('--num_groups', type=int, default=10, help="number of groups")
    parser.add_argument('--num_block', type=int, default=4, help="number of block")
    parser.add_argument('--group1_frac', type=float, default=0.5, help="the fraction of group1: A")
    parser.add_argument('--group2_frac', type=float, default=0.5, help="the fraction of group2: B")
    parser.add_argument('--num_id1',type=int, default=2 ,help = "number of id1")
    parser.add_argument('--num_id2',type=int, default=2 ,help = "number of id2")
    parser.add_argument('--client1_frac',type=float, default=0.4 ,help = "the fraction of client1")
    parser.add_argument('--client2_frac',type=float, default=0.4 ,help = "the fraction of client2")
    parser.add_argument('--client3_frac',type=float, default=0.2 ,help = "the fraction of client3")
    # parser.add_argument('--group1_frac', type=float, default=0.5, help="the fraction of group1: A")
    # parser.add_argument('--group2_frac', type=float, default=0.5, help="the fraction of group2: B")
    # parser.add_argument('--group3_frac', type=float, default=0.4, help="the fraction of group3: C")

    parser.add_argument('--sever_ep', type=int, default=10, help="the number of sever epochs: E")
    parser.add_argument('--sever_bs', type=int, default=50, help="sever batch size: B")

    # parser.add_argument('--num_layers', type=int, default=3, help="the number of layers")
    parser.add_argument('--max_norm', type=float, default=10, help="the maximum of norm")
    # loss
    # parser.add_argument('--alpha',type=float,default=0.9,help="percentage of cs_loss")
    parser.add_argument('--alpha',type=float,default=1.0,help="percentage of kl_loss")
    parser.add_argument('--temp',type=float,default=1.0,help="the temperature of KD")
    parser.add_argument('--K',type=int,default=3,help="the number of block per model")
    parser.add_argument('--sim_func',default='cka', choices=['cka', 'rbf_cka', 'lr'],help="the sim_func")

    parser.add_argument("--depth_saved", default=[0], type=int, nargs='*',
                        help="the index of network start channel slim.")  # vgg 采用4-6-8 resnet采用2-3-4 其中resnet是用block为最小单位
    parser.add_argument("--width_ration", default=[0.1, 0.2, 1.0], type=float, nargs='*',
                        help="the info of model ration and model type.") #heterofl[0.2, 0.5, 1.0] ScaleFLRes[0.3, 0.5, 1.0] ti[0.2, 0.4, 1.0] [0.15, 0.3, 1.0]
    parser.add_argument("--client_hetero_ration", default='4:4:2', type=str,
                        help="the info of model ration and model type.")
    parser.add_argument("--n_blocks", default=[2,3,4], type=float, nargs='*',
                        help="the num of blocks in DepthFL.")

    args = parser.parse_args()
    return args